import asyncio
from io import StringIO
from fastmcp import FastMCP
from fastmcp import Client
import requests
import json
from datetime import datetime, timedelta
import pandas as pd
from typing import List, Literal, Any
from pytrends.request import TrendReq

# 需要安装依赖：
# pip install lxml pytrends

mcp = FastMCP("demo.mcp")

# 从 Yahoo Finance 获取实时热门股票榜单（涨幅/跌幅/成交量）
@mcp.tool()
async def get_market_movers(
    category: Literal["gainers", "losers", "most-active"] = "most-active",
    count: int = 25,
    market_session: Literal["regular", "pre-market", "after-hours"] = "regular"
) -> str:
    """
    Retrieves market mover data (gainers, losers, most active) from Yahoo Finance.

    Input Parameters:
        category: Type of market movers to retrieve:
            - 'gainers': Top performing stocks by percentage gain.
            - 'losers': Worst performing stocks by percentage loss.
            - 'most-active': Stocks with highest trading volume.
        count: Number of results to return (1–100). Defaults to 25.
        market_session: Market session (only applies to 'most-active'):
            - 'regular': Regular trading hours
            - 'pre-market': Before market open
            - 'after-hours': After market close

    Output:
        A JSON string containing a list of market movers, each with:
            - symbol: Ticker symbol
            - name: Company name
            - price: Latest price
            - change: Absolute price change
            - percent_change: Percent change
            - volume: Trading volume
        Returns a JSON error object if the request or parsing fails.
    """
    base_urls = {
        "most-active": {
            "regular": "https://finance.yahoo.com/most-active",
            "pre-market": "https://finance.yahoo.com/markets/stocks/pre-market",
            "after-hours": "https://finance.yahoo.com/markets/stocks/after-hours",
        },
        "gainers": "https://finance.yahoo.com/gainers",
        "losers": "https://finance.yahoo.com/losers",
    }
    BROWSER_HEADERS = {
    "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/121.0.0.0 Safari/537.36"
    }

    count = min(max(count, 1), 100)

    try:
        if category == "most-active":
            url = base_urls["most-active"].get(market_session)
            if not url:
                return json.dumps({"error": f"Invalid market_session: {market_session}"})
        else:
            url = base_urls.get(category)
            if not url:
                return json.dumps({"error": f"Invalid category: {category}"})

        url = f"{url}?count={count}&offset=0"

        resp = requests.get(url, headers=BROWSER_HEADERS, timeout=15)
        if resp.status_code != 200:
            return json.dumps({"error": f"Failed to fetch data: HTTP {resp.status_code}"})

        tables = pd.read_html(StringIO(resp.text))
        if not tables or tables[0].empty:
            return json.dumps({"error": f"No table data found for category={category}"})

        df = tables[0].loc[:, ~tables[0].columns.str.contains('^Unnamed')]

        records = df.head(count).to_dict(orient="records")
        return json.dumps(records, indent=2, ensure_ascii=False)

    except Exception as e:
        return json.dumps({"error": f"Unexpected error: {str(e)}"})

# 从 CNN 官网的公开 API 获取 “恐惧与贪婪指数（Fear & Greed Index）” 及其相关市场情绪指标，并返回一个结构化的字典数据
@mcp.tool()
async def get_cnn_fear_greed_index(
    indicators: list[
        Literal[
            "fear_and_greed",
            "fear_and_greed_historical",
            "put_call_options",
            "market_volatility_vix",
            "market_volatility_vix_50",
            "junk_bond_demand",
            "safe_haven_demand"
        ]
    ] | None = None
) -> dict:
    """
    Fetch and return the CNN Fear & Greed Index data.

    This function retrieves the CNN Fear & Greed Index and related market sentiment indicators
    from CNN's public API. By default, it returns all indicators except the historical time series.
    If a list of indicators is provided, the output will be filtered to include only the selected ones.

    Parameters:
        indicators (Optional[list[FearGreedIndicator]]): 
            A list of indicator names to filter the results. 
            If None (default), all available indicators (excluding historical data) are returned.
            Available indicators:
                - "fear_and_greed"
                - "fear_and_greed_historical"
                - "put_call_options"
                - "market_volatility_vix"
                - "market_volatility_vix_50"
                - "junk_bond_demand"
                - "safe_haven_demand"

    Returns:
        dict: A dictionary containing the selected indicators and their current values. 
              Historical time series data and large "data" arrays are removed for cleaner output.
    """
    CNN_FEAR_GREED_URL = "https://production.dataviz.cnn.io/index/fearandgreed/graphdata"
    headers = {
        "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64)",
        "Accept": "application/json, text/plain, */*",
    }

    response = requests.get(CNN_FEAR_GREED_URL, headers=headers, timeout=15)
    response.raise_for_status()
    raw_data = response.json()

    if not raw_data:
        raise ValueError("Empty response data")

    # Remove historical time series data arrays
    result = {
        k: {inner_k: inner_v for inner_k, inner_v in v.items() if inner_k != "data"}
        if isinstance(v, dict) else v
        for k, v in raw_data.items()
        if k != "fear_and_greed_historical"
    }

    # Filter by indicators if specified
    if indicators:
        if invalid := set(indicators) - set(result.keys()):
            raise ValueError(f"Invalid indicators: {list(invalid)}. Available: {list(result.keys())}")
        result = {k: v for k, v in result.items() if k in indicators}

    return result

# 获取加密货币市场的恐惧与贪婪指数
@mcp.tool()
def get_crypto_fear_greed_index() -> dict:
    """
    Fetch the current Crypto Fear & Greed Index.

    This function retrieves the latest market sentiment index for cryptocurrencies
    from the Alternative.me public API. The Crypto Fear & Greed Index is a popular
    indicator that measures the overall sentiment of the crypto market, ranging from
    extreme fear to extreme greed.

    Source: https://alternative.me/crypto/fear-and-greed-index/

    Returns:
        dict: A dictionary containing:
            - value (str): The current index value (0-100).
            - classification (str): The sentiment classification (e.g., "Fear", "Greed").
            - timestamp (str): The Unix timestamp of the measurement.
    """
    CRYPTO_FEAR_GREED_URL = "https://api.alternative.me/fng/"
    headers = {
        "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64)",
        "Accept": "application/json",
    }

    try:
        response = requests.get(CRYPTO_FEAR_GREED_URL, headers=headers, timeout=15)
        response.raise_for_status()
        data = response.json()
    except requests.RequestException as e:
        raise ValueError(f"Failed to fetch Crypto Fear & Greed Index: {e}")
    except ValueError:
        raise ValueError("Failed to parse JSON response from Crypto Fear & Greed API")

    if "data" not in data or not data["data"]:
        raise ValueError("Invalid response format from Crypto Fear & Greed API")

    current_data = data["data"][0]

    return {
        "value": current_data.get("value"),
        "classification": current_data.get("value_classification"),
        "timestamp": current_data.get("timestamp"),
    }

# 获取 Google Trends 指定关键词的相对搜索热度时间序列
@mcp.tool()
def get_google_trends(
    keywords: List[str],
    period_days: int = 7
) -> str:
    """
    Fetch Google Trends interest data for given keywords over a specified period.

    This function retrieves the relative search interest (0–100) for one or more keywords
    using the pytrends library, which is an unofficial API for Google Trends.

    Args:
        keywords (List[str]): A list of keywords to query (e.g., ["Bitcoin", "Ethereum"]).
        period_days (int): Number of days to look back from today. Default is 7.

    Returns:
        str: A JSON-formatted string representing time series data, where each entry contains:
            - date (str): Date in ISO format (YYYY-MM-DD)
            - keyword1, keyword2, ... (int): Search interest scores for each keyword

    Raises:
        ValueError: If no data is returned from Google Trends or if keywords are invalid.
    """

    # Validate input
    if not keywords or not isinstance(keywords, list):
        raise ValueError("`keywords` must be a non-empty list of strings.")
    if period_days <= 0:
        raise ValueError("`period_days` must be a positive integer.")

    # Convert period_days into Google Trends timeframe format (e.g., "2025-10-08 2025-10-15")
    end_date = datetime.utcnow().date()
    start_date = end_date - timedelta(days=period_days)
    timeframe = f"{start_date} {end_date}"

    # Initialize pytrends client
    pytrends = TrendReq(hl="en-US", tz=360)

    # Build the request payload
    try:
        pytrends.build_payload(keywords, timeframe=timeframe)
        df = pytrends.interest_over_time()
    except Exception as e:
        raise ValueError(f"Failed to fetch Google Trends data: {e}")

    # Validate response
    if df.empty:
        raise ValueError("No data returned from Google Trends.")

    # Remove 'isPartial' column if present
    if "isPartial" in df.columns:
        df = df[~df["isPartial"]].drop(columns=["isPartial"])

    # Reset index to get 'date' as a column
    df_reset = df.reset_index()

    # Convert DataFrame to list of dicts
    records = df_reset.to_dict(orient="records")

    # Convert to JSON string
    return json.dumps(records, indent=2, ensure_ascii=False)

GET_MARKET_MOVERS_SCHEMA = {
    "type": "function",
    "function": {
        "name": "get_market_movers",
        "description": "Retrieve Yahoo Finance market movers (top gainers, losers, or most active stocks) for the selected session.",
        "parameters": {
            "type": "object",
            "properties": {
                "category": {
                    "type": "string",
                    "enum": ["gainers", "losers", "most-active"],
                    "description": "The type of market movers to retrieve: 'gainers', 'losers', or 'most-active'."
                },
                "count": {
                    "type": "integer",
                    "description": "Number of results to return (1–100). Default is 25."
                },
                "market_session": {
                    "type": "string",
                    "enum": ["regular", "pre-market", "after-hours"],
                    "description": "The market session to query (only applies to 'most-active')."
                }
            },
            "required": []
        }
    }
}

GET_CNN_FEAR_GREED_INDEX_SCHEMA = {
    "type": "function",
    "function": {
        "name": "get_cnn_fear_greed_index",
        "description": "Fetch CNN Fear & Greed Index and selected sentiment indicators from CNN's public API.",
        "parameters": {
            "type": "object",
            "properties": {
                "indicators": {
                    "type": "array",
                    "items": {
                        "type": "string",
                        "enum": [
                            "fear_and_greed",
                            "fear_and_greed_historical",
                            "put_call_options",
                            "market_volatility_vix",
                            "market_volatility_vix_50",
                            "junk_bond_demand",
                            "safe_haven_demand"
                        ]
                    },
                    "description": "Optional list of indicator names to filter. If not provided, all current indicators are returned."
                }
            },
            "required": []
        }
    }
}

GET_CRYPTO_FEAR_GREED_INDEX_SCHEMA = {
    "type": "function",
    "function": {
        "name": "get_crypto_fear_greed_index",
        "description": "Fetch the current Crypto Fear & Greed Index from the Alternative.me public API, including value, classification, and timestamp.",
        "parameters": {
            "type": "object",
            "properties": {},
            "required": []
        }
    }
}

GET_GOOGLE_TRENDS_SCHEMA = {
    "type": "function",
    "function": {
        "name": "get_google_trends",
        "description": "Fetch Google Trends search interest time series for one or more keywords over a specified period.",
        "parameters": {
            "type": "object",
            "properties": {
                "keywords": {
                    "type": "array",
                    "items": {
                        "type": "string"
                    },
                    "description": "A list of keywords to query (e.g., ['Bitcoin', 'Ethereum'])."
                },
                "period_days": {
                    "type": "integer",
                    "description": "Number of days to look back from today. Default is 7."
                }
            },
            "required": ["keywords"]
        }
    }
}

# --- 测试入口 ---
async def main():
    client = Client(mcp)
    async with client:
        tools = await client.list_tools()
        print("Available tools:", tools)

        result = await client.call_tool("get_market_movers", {
            "category": "gainers",
            "count": 10
        })
        print("Market Movers Result:", result)

        result = await client.call_tool("get_cnn_fear_greed_index", {
            "indicators": ["fear_and_greed", "market_volatility_vix"]
        })
        print("Market Movers Result:", result)

        result = await client.call_tool("get_crypto_fear_greed_index", {})
        print("Market Movers Result:", result)


if __name__ == "__main__":
    asyncio.run(main())